import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression

# 设置中文字体
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['axes.unicode_minus'] = False

# === Step 1: 读取数据 ===
file_path = '附件(Attachment).xlsx'
df = pd.read_excel(file_path, 2)

# 提取变量
t = df['时间 t (Time t)'].values.reshape(-1, 1)
F = df['主路3的车流量 (Traffic flow on the Main road 3)'].values.reshape(-1, 1)

# === Step 2: 自动枚举 t_c 以寻找最优 MSE ===
min_mse = float('inf')
best_tc = None
best_params = {}

# 合理设定拐点搜索范围，避免两端拟合不充分
for tc_index in range(5, len(t) - 5):
    tc = t[tc_index][0]

    # Step 2.1：拟合支路1 f1(t) = a1*t + b1（全时间段）
    reg_f1 = LinearRegression().fit(t, F)
    a1 = reg_f1.coef_[0].item()
    b1 = reg_f1.intercept_.item()
    f1 = a1 * t + b1

    # Step 2.2：推算支路2真实值 f2 = F - f1
    f2_true = F - f1

    # Step 2.3：拟合支路2（分段线性）
    t_left, f2_left = t[:tc_index+1], f2_true[:tc_index+1]
    t_right, f2_right = t[tc_index+1:], f2_true[tc_index+1:]

    reg_left = LinearRegression().fit(t_left, f2_left)
    reg_right = LinearRegression().fit(t_right, f2_right)

    a2 = reg_left.coef_[0].item()
    b2 = reg_left.intercept_.item()
    a3 = reg_right.coef_[0].item()
    b3 = reg_right.intercept_.item()

    # Step 2.4：合成预测值 f2_hat 和主路预测 F_hat
    f2_hat = np.where(t <= tc, a2 * t + b2, a3 * t + b3)
    F_hat = f1 + f2_hat
    mse = np.mean((F - F_hat)**2)

    if mse < min_mse:
        min_mse = mse
        best_tc = tc
        best_params = {
            'a1': a1, 'b1': b1,
            'a2': a2, 'b2': b2,
            'a3': a3, 'b3': b3
        }

# === Step 3: 输出结果 ===
print("✅ 拟合完成，自动搜索最优拐点")
print(f"最优拐点 t_c = {best_tc:.2f} 分钟")
print(f"最小 MSE = {min_mse:.4e}")
print(f"支路1函数: f1(t) = {best_params['a1']:.4f} * t + {best_params['b1']:.4f}")
print("支路2函数: f2(t) =")
print(f"    {best_params['a2']:.4f} * t + {best_params['b2']:.4f}，当 t ≤ {best_tc:.2f}")
print(f"    {best_params['a3']:.4f} * t + {best_params['b3']:.4f}，当 t >  {best_tc:.2f}")

# === Step 4: 可视化 ===
t = t.flatten()
f1 = best_params['a1'] * t + best_params['b1']
f2 = np.where(t <= best_tc,
              best_params['a2'] * t + best_params['b2'],
              best_params['a3'] * t + best_params['b3'])
F_hat = f1 + f2

plt.figure(figsize=(12, 6))
plt.plot(t, F, 'ko-', label='主路实测 F(t)')
plt.plot(t, F_hat, 'g--', label='模型预测 F_hat(t)')
plt.plot(t, f1, 'b:', label='支路1 f1(t)')
plt.plot(t, f2, 'r:', label='支路2 f2(t)')
plt.axvline(best_tc, color='gray', linestyle='--', label=f'最优拐点 t_c = {best_tc:.2f}')
plt.xlabel("时间 t / 分钟")
plt.ylabel("车流量 / 标准车当量")
plt.title("问题一：最优拐点分段拟合结果")
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
